from pulp import LpMaximize, LpProblem, LpVariable, lpSum, PULP_CBC_CMD
from common_import import *

# 创建问题
model = LpProblem(name="maximize-profit", sense=LpMaximize)

# 地块、作物、时间范围
lands = [
    "A1",
    "A2",
    "A3",
    "A4",
    "A5",
    "A6",
    "B1",
    "B2",
    "B3",
    "B4",
    "B5",
    "B6",
    "B7",
    "B8",
    "B9",
    "B10",
    "B11",
    "B12",
    "B13",
    "B14",
    "C1",
    "C2",
    "C3",
    "C4",
    "C5",
    "C6",
]
crops = range(1, 16)  # 作物编号
years = range(2024, 2031)  # 时间范围：2024到2030

# 面积、产量等数据
area_of_land = {
    "A1": 80,
    "A2": 55,
    "A3": 35,
    "A4": 72,
    "A5": 68,
    "A6": 55,
    "B1": 60,
    "B2": 46,
    "B3": 40,
    "B4": 28,
    "B5": 25,
    "B6": 86,
    "B7": 55,
    "B8": 44,
    "B9": 50,
    "B10": 25,
    "B11": 60,
    "B12": 45,
    "B13": 35,
    "B14": 20,
    "C1": 15,
    "C2": 13,
    "C3": 15,
    "C4": 18,
    "C5": 27,
    "C6": 20,
}

total_sales = {
    1: 57000,
    2: 21850,
    3: 22400,
    4: 33040,
    5: 9875,
    6: 170840,
    7: 132750,
    8: 71400,
    9: 30000,
    10: 12500,
    11: 1500,
    12: 35100,
    13: 36000,
    14: 14000,
    15: 10000,
}
x = LpVariable.dicts(
    "x",
    ((i, j, t) for i in lands for j in crops for t in years),
    lowBound=0,
    cat="Inte",
)
z = LpVariable.dicts(
    "z", ((i, j, t) for i in lands for j in crops for t in years), cat="Binary"
)

# 新增变量：表示每种作物每年超出的部分
overproduction = LpVariable.dicts(
    "overproduction",
    ((j, t) for j in crops for t in years),
    lowBound=0,
    cat="Continuous",
)

# 目标函数：最大化总利润，并惩罚超出部分
model += lpSum(
    mapping.get_profit_idtype(j, mapping.get_land_type(i), "单季")
    * x[i, j, t]
    * mapping.get_yield_idtype(j, mapping.get_land_type(i), "单季")
    for i in lands
    for j in crops
    for t in years
) - lpSum(
    overproduction[j, t] * mapping.get_priceperjin_idtype(j, "单季")
    for j in crops
    for t in years
)

# 约束条件1: 每个作物在所有地的一年总产量小于等于总销量 + 超出的部分
for j in crops:
    for t in years:
        total_yield = lpSum(
            mapping.get_yield_idtype(j, mapping.get_land_type(i), "单季") * x[i, j, t]
            for i in lands
        )
        model += total_yield <= total_sales[j] + overproduction[j, t]  # 允许超出部分
        model += overproduction[j, t] >= total_yield - total_sales[j]  # 计算超出的部分

# 约束条件1: 每个作物在所有地的总产量小于总销量
# for j in crops:
#     for t in years:
#         model += lpSum(yield_per_mu[j] * x[i, j, t] for i in lands) <= total_sales[j]

# 约束条件: 2023年的已知种植方案影响到2024年
known_planting_2023 = {
    "A1": 6,
    "A2": 7,
    "A3": 7,
    "A4": 1,
    "A5": 4,
    "A6": 8,
    "B1": 6,
    "B2": 2,
    "B3": 3,
    "B4": 4,
    "B5": 5,
    "B6": 8,
    "B7": 6,
    "B8": 8,
    "B9": 9,
    "B10": 10,
    "B11": 1,
    "B12": 7,
    "B13": 14,
    "B14": 15,
    "C1": 11,
    "C2": 12,
    "C3": 1,
    "C4": 13,
    "C5": 6,
    "C6": 3,
}

# 约束条件2: 每个地任意连续三年种植的豆类粮食面积大于等于该地的面积
bean_crops = [1, 2, 3, 4, 5]
for i in lands:
    # 考虑2023, 2024, 2025三个连续年份
    model += (
        lpSum(
            x[i, j, 2024]
            + x[i, j, 2025]
            + (100 if (i, j, 2023) in known_planting_2023 else 0)
            for j in bean_crops
        )
        >= area_of_land[i]
    )

    # 其他年份的连续三年
    for t in range(2024, 2029):
        model += (
            lpSum(x[i, j, t + k] for j in bean_crops for k in range(3))
            >= area_of_land[i]
        )
    pass

# 约束条件3: 每个地每年种植的作物总面积小于等于该地的面积
for i in lands:
    for t in years:
        model += lpSum(x[i, j, t] for j in crops) <= area_of_land[i]

# 约束条件4: 不能重茬种植
M = 100
for i in lands:
    for j in crops:
        for t in range(2024, 2030):
            model += x[i, j, t] <= M * z[i, j, t]
            model += z[i, j, t] + z[i, j, t + 1] <= 1

# 2024不能重茬种植
for i, j in known_planting_2023.items():
    model += z[i, j, 2024] == 0  # 2024年不能重茬种植
# 启用CBC求解器的详细输出
solver = PULP_CBC_CMD(msg=True, gapRel=0.05)

# 求解模型
model.solve(solver)

# 输出结果
# for i in lands:
#     for j in crops:
#         for t in years:
#             if x[i, j, t].value() > 0:
#                 print(f"Land {i}, Year {t}, Crop {j}: Area {x[i, j, t].value()} mu")
for t in years:
    for i in lands:
        for j in crops:
            if x[i, j, t].value() > 0:
                print(f"Year {t}, Land {i}, Crop {j}: Area {x[i, j, t].value()} mu")

# 初始化一个字典来存储每年的总利润
annual_profit = {t: 0 for t in years}

# 逐年计算总利润
for t in years:
    # 计算当年的总收益
    total_revenue = sum(
        mapping.get_profit_idtype(j, mapping.get_land_type(i), "单季")
        * x[i, j, t].value()
        * mapping.get_yield_idtype(j, mapping.get_land_type(i), "单季")
        for i in lands
        for j in crops
    )

    # 计算当年的惩罚项（超出部分）
    total_penalty = sum(
        overproduction[j, t].value() * mapping.get_priceperjin_idtype(j, "单季")
        for j in crops
    )

    # 计算每年的净利润
    annual_profit[t] = total_revenue - total_penalty

    # 输出每年的总利润
    print(f"Year {t} Total Profit: {annual_profit[t]}")

# 输出总利润
total_profit = sum(annual_profit.values())
print(f"Total Profit: {total_profit}")
